from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import json
from typing import List, Set, Dict, Any

class ResponseFormatterMiddleware(BaseHTTPMiddleware):
    """响应格式化中间件，统一API响应格式"""
    
    def __init__(self, app, exclude_paths: Set[str] = None, exclude_prefixes: List[str] = None):
        """
        初始化中间件
        
        Args:
            app: FastAPI应用
            exclude_paths: 完全排除的路径集合，这些路径的响应将保持原样
            exclude_prefixes: 排除的路径前缀列表，以这些前缀开头的路径将保持原样
        """
        super().__init__(app)
        self.exclude_paths = exclude_paths or {"/docs", "/redoc", "/openapi.json"}
        self.exclude_prefixes = exclude_prefixes or ["/docs/", "/redoc/", "/openapi/"]
    
    async def dispatch(self, request: Request, call_next):
        # 调用下一个中间件或路由处理器
        response = await call_next(request)
        
        # 检查路径是否需要排除
        path = request.url.path
        if path in self.exclude_paths or any(path.startswith(prefix) for prefix in self.exclude_prefixes):
            return response
        
        # 检查响应头是否包含跳过标记
        if response.headers.get("X-Skip-Formatter", "").lower() == "true":
            # 移除头部后返回
            headers = dict(response.headers)
            headers.pop("X-Skip-Formatter", None)
            return response.__class__(
                content=response.body,
                status_code=response.status_code,
                headers=headers,
                media_type=response.media_type
            )
        
        # 只处理JSONResponse
        if isinstance(response, JSONResponse):
            try:
                # 解析原始响应内容
                body = response.body.decode()
                content = json.loads(body)
                
                # 检查是否已经是标准格式
                if isinstance(content, dict) and all(k in content for k in ["code", "data", "message"]):
                    return response
                
                # 构造新的标准格式响应
                status_code = response.status_code
                new_content = {
                    "code": str(status_code),
                    "data": content,
                    "message": "success" if status_code < 400 else "error"
                }
                
                # 创建新的JSONResponse，保留原始状态码和头信息
                headers = dict(response.headers)
                return JSONResponse(
                    content=new_content,
                    status_code=status_code,
                    headers=headers
                )
            except json.JSONDecodeError:
                # 如果解析失败，返回原始响应
                return response
        
        # 对于非JSON响应，保持原样
        return response

def create_standard_response(
    data: Any = None, 
    code: str = "200", 
    message: str = "success"
) -> Dict[str, Any]:
    """
    创建标准格式的响应
    
    Args:
        data: 响应数据
        code: 状态码
        message: 响应消息
        
    Returns:
        标准格式的响应字典
    """
    return {
        "code": code,
        "data": data,
        "message": message
    }

# 创建跳过格式化的响应
def skip_formatter(response: JSONResponse) -> JSONResponse:
    """
    标记响应跳过格式化
    
    Args:
        response: 原始响应对象
        
    Returns:
        带有跳过格式化标记的响应对象
    """
    response.headers["X-Skip-Formatter"] = "true"
    return response 